import os, re, json, requests
from typing import List, Dict, Any
from dotenv import load_dotenv
from typing import Optional, Union

load_dotenv()

class LLMClient:
    def __init__(self):
        self.provider = os.getenv("LLM_PROVIDER", "ollama")
        self.host = os.getenv("LLM_HOST", "http://localhost:11434")
        self.model = os.getenv("LLM_MODEL", "deepseek-r1:1.5b")
        self.api_key = os.getenv("LLM_API_KEY", "")

    def generate(self, prompt: str) -> str:
        if self.provider == "ollama":
            resp = requests.post(f"{self.host}/api/generate", json={
                "model": self.model, "prompt": prompt, "stream": False
            })
            resp.raise_for_status()
            return resp.json().get("response", "").strip()

        else:
            raise ValueError(f"Unsupported LLM provider: {self.provider}")

class DataCleaningAgent:
    def __init__(self, llm: Optional[LLMClient] = None):
        self.llm = llm or LLMClient()

    @staticmethod
    def clean_text(text: str) -> str:
        text = re.sub(r"<[^>]+>", " ", text) # 去HTML
        text = re.sub(r"\s+", " ", text) # 合并空白
        text = re.sub(r"[^\w\s\u4e00-\u9fa5@：:，,。.\-—/()（）]+", " ", text) # 去特殊符
        return text.strip()

    def extract_fields(self, text: str, instruction: str) -> Dict[str, Any]:
        """
        instruction 示例：
        从文本中提取公司名称(company)、联系人(contact)、电话(phone)。返回JSON，如：
        {"company": "...", "contact": "...", "phone": "..."}
        """
        prompt = f"""你是信息抽取助手。请严格输出JSON。

        文本：{text}

        任务：{instruction}
        
        注意： 
        1. 只输出一个JSON对象，不要输出多余文字。
        2. 缺失字段请用null。
        """

        raw = self.llm.generate(prompt)

        return parse_llm_json_response(raw)
        # try:
        #     return json.loads(raw)
        # except Exception:
        #     # LLM 偶发输出非纯JSON，兜底提取
        #     m = re.search(r"\{.*\}", raw, re.S)
        #     return json.loads(m.group(0)) if m else {"_raw": raw}

    def batch_process(self, samples: List[str], instruction: str) -> List[Dict[str,Any]]:
        out = []
        for s in samples:
            cleaned = self.clean_text(s)
            extracted = self.extract_fields(cleaned, instruction)
            out.append({"original": s, "cleaned": cleaned, "extracted": extracted})
        return out
    

def parse_llm_json_response(raw: str) -> Dict[str, Any]:
    """
    解析 LLM 生成的 JSON 响应，带有强大的错误处理和回退机制
    
    Args:
        raw: LLM 生成的原始响应字符串
        
    Returns:
        解析后的 JSON 字典，或包含原始响应的字典
    """
    # 1. 尝试直接解析
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        pass
    
    # 2. 清理字符串中的常见问题
    cleaned = raw.strip()
    
    # 移除常见的代码块标记
    cleaned = re.sub(r'^```(?:json)?\s*|\s*```$', '', cleaned, flags=re.MULTILINE)
    
    # 移除JSON中的单行注释
    cleaned = re.sub(r'^\s*//.*$', '', cleaned, flags=re.MULTILINE)
    
    # 3. 尝试解析清理后的内容
    try:
        return json.loads(cleaned)
    except json.JSONDecodeError:
        pass
    
    # 4. 使用更精确的JSON提取（支持嵌套结构）
    def find_json_substring(text: str):
        """使用栈匹配找到最外层完整JSON对象"""
        start = text.find('{')
        if start == -1:
            start = text.find('[')  # 也支持数组
            if start == -1:
                return None
        
        # 确定是对象还是数组
        stack = []
        for i in range(start, len(text)):
            char = text[i]
            if char in '{[':
                stack.append(char)
            elif char in '}]':
                if not stack:
                    break
                opening = stack.pop()
                if (opening == '{' and char != '}') or (opening == '[' and char != ']'):
                    break
            
            # 如果栈为空，找到了完整的JSON
            if not stack:
                # 向前查找可能更早的结束位置
                for j in range(i + 1, len(text)):
                    if text[j] in '{[':
                        # 找到新的开始，提前结束当前JSON
                        i = j - 1
                        break
                return text[start:i+1]
        return None
    
    # 5. 尝试找到并解析JSON子串
    json_str = find_json_substring(cleaned)
    if json_str:
        try:
            return json.loads(json_str)
        except json.JSONDecodeError:
            # 尝试更多清理
            json_str = clean_json_string(json_str)
            try:
                return json.loads(json_str)
            except json.JSONDecodeError:
                pass
    
    # 6. 最后的尝试：使用更宽松的解析
    try:
        # 尝试解析为Python字面量（可能包含单引号）
        import ast
        result = ast.literal_eval(cleaned.replace('\n', ' ').replace('\r', ''))
        if isinstance(result, (dict, list)):
            return result if isinstance(result, dict) else {"_array": result}
    except (SyntaxError, ValueError):
        pass
    
    # 所有尝试都失败，返回原始文本
    return {"_raw": raw[:500], "_error": "Failed to parse JSON response", "_truncated": len(raw) > 500}

def clean_json_string(json_str: str) -> str:
    """清理JSON字符串中的常见问题"""
    if not json_str:
        return json_str
    
    # 1. 修复未转义的控制字符
    json_str = re.sub(r'(?<!\\)\\(?!["\\/bfnrtu])', r'\\\\', json_str)
    
    # 2. 修复单引号（有时LLM会用单引号）
    # 首先处理字符串内的单引号
    def replace_single_quotes(match):
        content = match.group(1)
        # 转义内容中的双引号
        content = content.replace('"', '\\"')
        return f'"{content}"'
    
    # 匹配单引号字符串，忽略转义的单引号
    pattern = r"'((?:[^'\\]|\\.)*)'"
    json_str = re.sub(pattern, replace_single_quotes, json_str)
    
    # 3. 修复未转义换行符
    json_str = json_str.replace('\n', '\\n').replace('\r', '\\r')
    
    # 4. 修复尾随逗号
    json_str = re.sub(r',\s*([}\]])', r'\1', json_str)
    
    # 5. 移除JavaScript风格的注释（单行和多行）
    json_str = re.sub(r'//.*?(\n|$)', '', json_str)
    json_str = re.sub(r'/\*.*?\*/', '', json_str, flags=re.DOTALL)
    
    return json_str.strip()

